import time

import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt
import numpy as np
import control
from mpl_toolkits.mplot3d import Axes3D


class Quadrotor:
    def __init__(self, mass, Ixx, Iyy, Izz, g,dt):
        # 物理参数
        self.mass = mass # 质量 (kg)
        self.g = g # 重力加速度 (m/s²)
        self.Ixx =  Ixx  # 绕x轴转动惯量 (kg·m²)
        self.Iyy = Iyy  # 绕y轴转动惯量 (kg·m²)
        self.Izz = Izz
        self.dt = dt  # 时间步长 (s)

        # 状态变量：[x, y, z, vx, vy, vz, yaw, pitch,roll, p, q, r]
        self.n_states = 12
        self.n_controls = 4  # [F,M1,M2,M3]

        # 状态空间模型
        self.A,self.B = self.get_dLTI()

    def step(self, state, control):
        return self.A @ state + self.B @ control

    def get_dLTI(self):
        num_x = 12
        Ac = np.array([[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, -self.g, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, self.g, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], ])
        Bc = np.array([[0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [-1 / self.mass, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 1/self.Ixx,0, 0],
                       [0, 0, 1/self.Iyy, 0],
                       [0, 0, 0, 1/self.Izz]])
        Cc = np.array([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0]])
        Dc = np.zeros((6, 4))
        sysc = control.ss(Ac, Bc, Cc, Dc)
        # Discretization 双线性变换离散化
        sysd = control.sample_system(sysc, self.dt, method='bilinear')
        print("sysd.A", sysd.A)
        print("sysd.B", sysd.B)
        print("Bc", Bc)
        return sysd.A, sysd.B

class TrajectoryGenerator:
    def __init__(self, total_time):
        self.total_time = total_time

    def get_desired_state(self, t):
        # 生成螺旋轨迹（三维）
        radius = 2.0
        omega = 0.2
        z_amplitude = 0.5
        x_d = radius * np.cos(omega * t)
        y_d = radius * np.sin(omega * t)
        z_d = z_amplitude * np.sin(omega * t)
        vx_d = -radius * omega * np.sin(omega * t)
        vy_d = radius * omega * np.cos(omega * t)
        vz_d = z_amplitude * omega * np.cos(omega * t)

        # 计算角度和角速度
        roll_d = 0.0
        pitch_d = 0.0
        yaw_d = omega * t + np.pi / 2  # 保持航向指向切线方向
        p_d = 0.0  # roll角速度
        q_d = 0.0  # pitch角速度
        r_d = omega  # yaw角速度

        # 合并所有状态
        return np.array([x_d, y_d, z_d, vx_d, vy_d, vz_d, roll_d, pitch_d, yaw_d, p_d, q_d, r_d])


class MPCController:
    def __init__(self, quadrotor, N, Q, R):
        self.quadrotor = quadrotor
        self.N = N  # 预测时域
        self.Q = Q  # 状态权重矩阵（6x6）
        self.R = R  # 控制输入权重矩阵（4x4）

    def control(self, state, desired_states):
        n_states = self.quadrotor.n_states
        n_controls = self.quadrotor.n_controls
        x = cp.Variable((n_states, self.N + 1))
        u = cp.Variable((n_controls, self.N))

        cost = 0
        constraints = []

        # 初始状态约束
        constraints.append(x[:, 0] == state)

        for k in range(self.N):
            # 状态转移约束
            constraints.append(
                x[:, k + 1] ==
                self.quadrotor.A @ x[:, k] +
                self.quadrotor.B @ u[:, k]
            )

            # 代价函数
            cost += cp.quad_form(x[:, k] - desired_states[:, k], self.Q)
            cost += cp.quad_form(u[:, k], self.R)

        # 终端代价
        cost += cp.quad_form(x[:, -1] - desired_states[:, -1], self.Q)

        prob = cp.Problem(cp.Minimize(cost), constraints)
        prob.solve(solver=cp.OSQP)

        return u[:, 0].value


def simulate():
    quadrotor = Quadrotor(mass=1.0, Ixx=0.00089, Iyy=0.00221, Izz=0.00289, g=9.8,dt=0.2)
    quadrotor.get_dLTI()
    # traj_gen = TrajectoryGenerator(total_time=20)
    # N = 10 # 预测时域
    # Q = np.diag([1.0, 1.0, 1.0, 0.1, 0.1, 0.1,1.0, 1.0, 1.0, 0.1, 0.1, 0.1])  # 状态权重（位置权重高）
    # R = np.diag([0.1, 0.01, 0.01, 0.01])  # 输入权重（总力权重较高）
    # mpc = MPCController(quadrotor, N, Q, R)
    #
    # # 初始化状态（位置0，速度0）
    # state = np.zeros(quadrotor.n_states)
    # actual_trajectory = []
    # desired_trajectory = []
    #
    # total_time = 20
    #
    # num_steps = int(total_time / quadrotor.dt)
    # time_steps = np.linspace(0, total_time, num_steps)
    # for i in range(num_steps):
    #     start_time=time.time()*1000
    #     t = i * quadrotor.dt
    #     # 生成未来N步的期望状态
    #     desired_states = np.zeros((quadrotor.n_states, N + 1))
    #     for k in range(N + 1):
    #         desired_states[:, k] = traj_gen.get_desired_state(t + k * quadrotor.dt)
    #
    #     # 计算控制输入
    #     control = mpc.control(state, desired_states)
    #
    #     # 更新状态
    #     state = quadrotor.step(state, control)
    #
    #     # 存储轨迹
    #     actual_trajectory.append(state[:3])
    #     desired_trajectory.append(desired_states[:3, 0])
    #     end_time = time.time()*1000
    #     print("耗时:",end_time - start_time)
    # # 绘制三维轨迹
    # fig = plt.figure(figsize=(12, 6))
    # ax = fig.add_subplot(111, projection='3d')
    # actual = np.array(actual_trajectory)
    # desired = np.array(desired_trajectory)
    #
    # ax.plot(desired[:, 0], desired[:, 1], desired[:, 2], 'b--', label='Desired')
    # ax.plot(actual[:, 0], actual[:, 1], actual[:, 2], 'r-', label='Actual')
    # ax.set_xlabel('X (m)'), ax.set_ylabel('Y (m)'), ax.set_zlabel('Z (m)')
    # ax.legend()
    # ax.set_title('3D Trajectory Tracking')
    # plt.figure()
    # plt.plot(time_steps, actual[:,0], 'r-', label='Actual')
    # plt.plot(time_steps, desired[:,0], 'b--', label='Desired')
    # plt.xlabel('time (s)'), plt.ylabel('X (m)')
    # plt.legend()
    # plt.title('X Trajectory Tracking')
    # plt.figure()
    # plt.plot(time_steps, actual[:, 1], 'r-', label='Actual')
    # plt.plot(time_steps, desired[:, 1], 'b--', label='Desired')
    # plt.xlabel('time (s)'), plt.ylabel('Y (m)')
    # plt.legend()
    # plt.title('Y Trajectory Tracking')
    # plt.figure()
    # plt.plot(time_steps, actual[:, 2], 'r-', label='Actual')
    # plt.plot(time_steps, desired[:, 2], 'b--', label='Desired')
    # plt.xlabel('time (s)'), plt.ylabel('Z (m)')
    # plt.legend()
    # plt.title('Z Trajectory Tracking')
    # plt.tight_layout()
    # plt.show()


if __name__ == "__main__":
    simulate()
